#!/usr/bin/env python3
"""
Complete system validation script for Manga Colorizer
Tests all components: API, models, storage, and deployment
"""

import requests
import json
import time
import sys
import os
from typing import Dict, Any, List
import subprocess
import tempfile
from PIL import Image
import io

# Configuration
API_BASE_URL = "http://localhost:8000/api/v1"
MAX_RETRIES = 30
RETRY_DELAY = 2


class SystemValidator:
    """Comprehensive system validation"""
    
    def __init__(self):
        self.results = []
        self.passed = 0
        self.total = 0
    
    def log_result(self, test_name: str, passed: bool, message: str = ""):
        """Log test result"""
        self.total += 1
        if passed:
            self.passed += 1
            print(f"✅ {test_name}: {message}")
        else:
            print(f"❌ {test_name}: {message}")
        self.results.append({
            "test": test_name,
            "passed": passed,
            "message": message
        })
    
    def test_docker_services(self) -> bool:
        """Test if Docker services are running"""
        try:
            # Check if services are running
            result = subprocess.run(
                ["docker-compose", "ps", "--services", "--filter", "status=running"],
                capture_output=True,
                text=True,
                cwd=os.path.dirname(os.path.dirname(__file__))
            )
            running_services = result.stdout.strip().split('\n')
            expected_services = ['api', 'redis', 'minio']
            
            for service in expected_services:
                if service not in running_services:
                    self.log_result(
                        f"Docker Service {service}",
                        False,
                        f"Service {service} is not running"
                    )
                    return False
            
            self.log_result("Docker Services", True, "All required services are running")
            return True
            
        except Exception as e:
            self.log_result("Docker Services", False, str(e))
            return False
    
    def test_api_health(self) -> bool:
        """Test API health endpoint"""
        try:
            response = requests.get(f"{API_BASE_URL}/health", timeout=10)
            if response.status_code == 200:
                data = response.json()
                if data.get('status') == 'healthy':
                    self.log_result("API Health", True, "API is healthy")
                    return True
                else:
                    self.log_result("API Health", False, f"API reports unhealthy: {data}")
                    return False
            else:
                self.log_result("API Health", False, f"Status code: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("API Health", False, str(e))
            return False
    
    def test_api_endpoints(self) -> bool:
        """Test all API endpoints"""
        endpoints = [
            ("GET", f"{API_BASE_URL}/health"),
            ("GET", f"{API_BASE_URL}/colorize/status/test"),
            ("GET", f"{API_BASE_URL}/colorize/result/test"),
            ("GET", f"{API_BASE_URL}/colorize/tasks"),
        ]
        
        all_passed = True
        for method, url in endpoints:
            try:
                if method == "GET":
                    response = requests.get(url, timeout=5)
                    # 404 is acceptable for test IDs
                    if response.status_code in [200, 404]:
                        self.log_result(f"API {method} {url}", True, "Endpoint accessible")
                    else:
                        self.log_result(f"API {method} {url}", False, f"Status: {response.status_code}")
                        all_passed = False
            except Exception as e:
                self.log_result(f"API {method} {url}", False, str(e))
                all_passed = False
        
        return all_passed
    
    def create_test_image(self) -> bytes:
        """Create a simple test image"""
        img = Image.new('RGB', (512, 512), color='white')
        
        # Add some simple lines to simulate a sketch
        from PIL import ImageDraw
        draw = ImageDraw.Draw(img)
        
        # Draw a simple manga-style face
        draw.ellipse([(150, 100), (350, 300)], outline='black', width=3)
        draw.ellipse([(200, 150), (220, 170)], outline='black', width=2)
        draw.ellipse([(280, 150), (300, 170)], outline='black', width=2)
        draw.line([(200, 220), (300, 220)], fill='black', width=2)
        
        buffer = io.BytesIO()
        img.save(buffer, format='PNG')
        return buffer.getvalue()
    
    def test_colorization_flow(self) -> bool:
        """Test complete colorization flow"""
        try:
            # Create test image
            test_image = self.create_test_image()
            
            # Upload and colorize
            files = {'file': ('test_sketch.png', test_image, 'image/png')}
            data = {
                'style': 'anime',
                'strength': 0.7,
                'guidance_scale': 7.5,
                'num_inference_steps': 5  # Reduced for testing
            }
            
            response = requests.post(
                f"{API_BASE_URL}/colorize",
                files=files,
                data=data,
                timeout=30
            )
            
            if response.status_code != 200:
                self.log_result("Colorization Upload", False, f"Status: {response.status_code}")
                return False
            
            result = response.json()
            task_id = result.get('task_id')
            
            if not task_id:
                self.log_result("Colorization Upload", False, "No task ID returned")
                return False
            
            self.log_result("Colorization Upload", True, f"Task ID: {task_id}")
            
            # Poll for completion
            for attempt in range(MAX_RETRIES):
                status_response = requests.get(
                    f"{API_BASE_URL}/colorize/status/{task_id}",
                    timeout=10
                )
                
                if status_response.status_code == 200:
                    status = status_response.json()
                    
                    if status['status'] == 'completed':
                        self.log_result("Colorization Complete", True, f"Completed in {attempt + 1} polls")
                        return True
                    elif status['status'] == 'failed':
                        self.log_result("Colorization Failed", False, status.get('error_message', 'Unknown error'))
                        return False
                
                time.sleep(RETRY_DELAY)
            
            self.log_result("Colorization Timeout", False, f"Timeout after {MAX_RETRIES} attempts")
            return False
            
        except Exception as e:
            self.log_result("Colorization Flow", False, str(e))
            return False
    
    def test_storage_service(self) -> bool:
        """Test storage service"""
        try:
            # Test MinIO/S3 connectivity
            response = requests.get("http://localhost:9000/minio/health/live", timeout=5)
            if response.status_code == 200:
                self.log_result("MinIO Storage", True, "Storage service is healthy")
                return True
            else:
                self.log_result("MinIO Storage", False, f"Status: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("MinIO Storage", False, str(e))
            return False
    
    def test_redis_service(self) -> bool:
        """Test Redis service"""
        try:
            import redis
            r = redis.from_url("redis://localhost:6379/0")
            r.ping()
            self.log_result("Redis Service", True, "Redis is accessible")
            return True
        except Exception as e:
            self.log_result("Redis Service", False, str(e))
            return False
    
    def test_model_loading(self) -> bool:
        """Test AI model loading"""
        try:
            from app.services.colorizer import colorization_service
            
            if colorization_service.is_model_loaded():
                self.log_result("AI Model Loading", True, "Model is loaded and ready")
                return True
            else:
                self.log_result("AI Model Loading", False, "Model not loaded")
                return False
        except Exception as e:
            self.log_result("AI Model Loading", False, str(e))
            return False
    
    def test_configuration(self) -> bool:
        """Test configuration"""
        try:
            from app.core.config import settings
            
            required_vars = [
                'model_name',
                'controlnet_model',
                'max_file_size',
                'allowed_extensions'
            ]
            
            missing_vars = []
            for var in required_vars:
                if not hasattr(settings, var) or getattr(settings, var) is None:
                    missing_vars.append(var)
            
            if missing_vars:
                self.log_result("Configuration", False, f"Missing: {missing_vars}")
                return False
            
            self.log_result("Configuration", True, "All required settings present")
            return True
            
        except Exception as e:
            self.log_result("Configuration", False, str(e))
            return False
    
    def run_all_tests(self) -> Dict[str, Any]:
        """Run all validation tests"""
        print("🧪 Starting Manga Colorizer System Validation")
        print("=" * 50)
        
        tests = [
            ("Docker Services", self.test_docker_services),
            ("API Health", self.test_api_health),
            ("API Endpoints", self.test_api_endpoints),
            ("Storage Service", self.test_storage_service),
            ("Redis Service", self.test_redis_service),
            ("Configuration", self.test_configuration),
            ("Colorization Flow", self.test_colorization_flow),
        ]
        
        for test_name, test_func in tests:
            print(f"\n{test_name}:")
            test_func()
        
        print("\n" + "=" * 50)
        print(f"Validation Results: {self.passed}/{self.total} tests passed")
        
        if self.passed == self.total:
            print("🎉 All tests passed! System is ready for production use.")
        else:
            print("❌ Some tests failed. Check the logs above.")
        
        return {
            "total_tests": self.total,
            "passed_tests": self.passed,
            "failed_tests": self.total - self.passed,
            "results": self.results,
            "success": self.passed == self.total
        }


def main():
    """Main validation function"""
    validator = SystemValidator()
    results = validator.run_all_tests()
    
    # Exit with appropriate code
    sys.exit(0 if results["success"] else 1)


if __name__ == "__main__":
    main()
